#Written by Felix Hsiao
#Last edit on 8/23/2013 by Felix Hsiao

import time

class prologix_33250a:
    def __init__(self, prologix, addr, function=None, freq=None, ampl=None, offset=None, load='INF', debug=False):
        self.prologix = prologix
        self.addr = addr
        self.function = function  #Default sine
        self.freq = freq  #Default 1 kHz
        self.ampl = ampl  #Default 0.2 Vpp (high-Z); can specify rms
        self.offset = offset  #Default 0
        self.load = load  #Output termination
        self.debug = debug
        self.initialize()

    def initialize(self):
        if self.debug: print 'Initializing address %d for outputting %s ...' %(self.addr, self.function)
        self.prologix.set_address(self.addr)
        if self.debug: print 'Reset'
        self.prologix.write('*RST')
        if self.debug: print 'Clear'
        self.prologix.write('*CLS')
        if self.debug: print 'Set function'
        if self.function is not None:
            self.prologix.write('FUNC %s' %(self.function))
        if self.debug: print 'Set output termination'
        self.prologix.write('OUTP:LOAD %s' %(self.load))
        if self.freq is not None:
            if self.debug: print 'Set frequency'
            self.prologix.write('FREQ %s' %(str(self.freq)))
        if self.ampl is not None:
            if self.debug: print 'Set amplitude'
            self.prologix.write('VOLT %s' %(str(self.ampl)))
        if self.offset:
            if self.debug: print 'Set offset'
            self.prologix.write('VOLT:OFFS %s' %(str(self.offset)))
        if self.debug: print 'Setup triggering'
        self.prologix.write('TRIG:SOUR BUS')  #Only used for burst and sweep
        if self.debug: print 'Address %d initialized.' %(self.addr)

    def activate(self):
        if self.debug: print 'Starting output from address %d ...' %(self.addr)
        self.prologix.set_address(self.addr)
        if self.debug: print 'OUTP ON'
        self.prologix.write('OUTP ON')
        if self.debug: print 'Address %d output on.' %(self.addr)
    
    def deactivate(self):
        if self.debug: print 'Stopping output from address %d ...' %(self.addr)
        self.prologix.set_address(self.addr)
        if self.debug: print 'OUTP OFF'
        self.prologix.write('OUTP OFF')
        if self.debug: print 'Address %d output off.' %(self.addr)

    def configSweep(self, freq1, freq2, centered=False, time=None, type=None):
        if self.debug: print 'Setting up frequency sweep for address %d ...' %(self.addr)
        self.prologix.set_address(self.addr)
        if not centered:  #freq1 and freq2 are start and stop points
            if self.debug: print 'Set frequency endpoints'
            self.prologix.write('FREQ:STAR %s' %(str(freq1)))
            self.prologix.write('FREQ:STOP %s' %(str(freq2)))
        else:  #freq1 is center frequency, freq2 is total span
            if self.debug: print 'Set center frequency and span'
            self.prologix.write('FREQ:CENT %s' %(str(freq1)))
            self.prologix.write('FREQ:SPAN %s' %(str(freq2)))
        if time is not None:  #Default 1 second
            if self.debug: print 'Set sweep time'
            self.prologix.write('SWE:TIME %s' %(str(time)))
        if type is not None:  #Can be LINear or LOGarithmic, default linear
            if self.debug: print 'Set sweep type'
            self.prologix.write('SWE:SPAC %s' %(type))
        if self.debug: print 'SWE ON'
        self.prologix.write('SWE:STAT ON')
        if self.debug: print 'Address %d sweep mode configured.' %(self.addr)

    def configBurst(self, cycles):
        if self.debug: print 'Setting up burst mode for address %d ...' %(self.addr)
        self.prologix.set_address(self.addr)
        if self.debug: print 'Set cycles'
        self.prologix.write('BURS:NCYC %s' %(str(cycles)))
        if self.debug: print 'BURS ON'
        self.prologix.write('BURS:STAT ON')
        if self.debug: print 'Address %d burst mode configured.' %(self.addr)
    
    def checkError(self):
        if self.debug: print 'Checking if error was generated by address %d ...' %(self.addr)
        self.prologix.set_address(self.addr)
        if self.debug: print 'ERR?'
        self.prologix.write('SYST:ERR?')
        time.sleep(0.1)
        data = self.prologix.readline()
        if self.debug: print 'Recv:', data
        return data
